#!pip install pypdf
#!pip install langchain_text_splitters
#!pip install faiss-cpu
#!pip install langchain_community

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.llms import Tongyi
from langchain.chains import ConversationalRetrievalChain
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from dotenv import load_dotenv
import os

# 加载环境变量
load_dotenv()

def main():
    print("=== RAG PDF 问答系统 ===")
    
    # 1. 从PDF文件读取内容
    print("\n1. 读取PDF文件...")
    file_path = 'G:\\TestProgram\\Langchain_study\\pdf_file\\论文介绍.pdf'
    loader = PyPDFLoader(file_path)
    pages = loader.load()
    print(f"成功加载PDF，共{len(pages)}页")
    
    # 2. 文本分割
    print("\n2. 分割文本...")
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,  # 每个文本块的最大字符数
        chunk_overlap=50,  # 相邻文本块之间的重叠字符数
        separators=["\n", "。", "！", "？", "，", "、", ""]  # 文本分割符
    )
    docs = text_splitter.split_documents(pages)
    print(f"文本分割完成，共{len(docs)}个文档块")
    
    # 3. 使用通义的嵌入模型
    print("\n3. 初始化嵌入模型...")
    api_key = os.getenv("DASHSCOPE_API_KEY")
    embeddings = DashScopeEmbeddings(
        model="text-embedding-v1",
        dashscope_api_key=api_key
    )
    print("嵌入模型初始化完成")
    
    # 4. 使用FAISS创建向量数据库
    print("\n4. 创建FAISS向量数据库...")
    vectorstore = FAISS.from_documents(docs, embeddings)
    print("向量数据库创建完成")
    
    # 5. 创建检索器
    print("\n5. 创建检索器...")
    retriever = vectorstore.as_retriever(
        search_type="similarity",
        search_kwargs={"k": 3}  # 返回最相似的3个文档块
    )
    print("检索器创建完成")
    
    # 6. 创建一个带记忆的对话检索链，使用通义模型
    print("\n6. 创建对话检索链...")
    llm = Tongyi(
        model="qwen-turbo",
        api_key=api_key,
        temperature=0.7,
        max_tokens=1024
    )
    
    # 创建基础检索链（不带记忆）
    base_qa_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=retriever,
        return_source_documents=True
    )
    
    # 创建内存存储
    store = {}
    
    def get_session_history(session_id: str) -> InMemoryChatMessageHistory:
        if session_id not in store:
            store[session_id] = InMemoryChatMessageHistory()
        return store[session_id]
    
    # 创建带历史记录的可运行对象
    qa_chain = RunnableWithMessageHistory(
        base_qa_chain,
        get_session_history,
        input_messages_key="question",
        history_messages_key="chat_history",
    )
    print("对话检索链创建完成")
    
    # 7. 测试调用链，传入用户提问
    print("\n7. 测试问答系统...")
    
    # 设置会话ID
    session_id = "rag_user_123"
    
    # 第一个问题
    question1 = "这篇论文的主要内容是什么？"
    print(f"\n问题1: {question1}")
    result1 = qa_chain.invoke(
        {"question": question1},
        config={"configurable": {"session_id": session_id}}
    )
    
    # 8. 打印模型响应的结果
    print("\n=== 回答1 ===")
    print(result1["answer"])
    print(f"\n参考文档数量: {len(result1['source_documents'])}")
    
    # 第二个问题（测试记忆功能）
    question2 = "论文中提到的主要方法有哪些？"
    print(f"\n问题2: {question2}")
    result2 = qa_chain.invoke(
        {"question": question2},
        config={"configurable": {"session_id": session_id}}
    )
    
    print("\n=== 回答2 ===")
    print(result2["answer"])
    print(f"\n参考文档数量: {len(result2['source_documents'])}")
    
    # 第三个问题（测试上下文理解）
    question3 = "能详细解释一下刚才提到的方法吗？"
    print(f"\n问题3: {question3}")
    result3 = qa_chain.invoke(
        {"question": question3},
        config={"configurable": {"session_id": session_id}}
    )
    
    print("\n=== 回答3 ===")
    print(result3["answer"])
    
    print("\n=== RAG系统测试完成 ===")
    
    # 可选：显示检索到的源文档
    print("\n=== 最后一次检索的源文档片段 ===")
    for i, doc in enumerate(result3['source_documents'][:2]):
        print(f"\n文档片段 {i+1}:")
        print(doc.page_content[:200] + "...")

if __name__ == "__main__":
    main()

